# 检查依赖库是否安装
try:
    import sys
    import platform
    import subprocess
    import urllib.request
    import tempfile
    import os
except ImportError as e:
    print(f"错误: 缺少必要的依赖库 - {e.name}")
    print("请先安装依赖库:")
    print(f"  {sys.executable} -m pip install {e.name}")
    sys.exit(1)


# 版本映射表
VERSION_MAP = {
    "3.10": "cp310-cp310",
    "3.11": "cp311-cp311", 
    "3.12": "cp312-ccp312"
}

def get_python_version():
    """获取Python主次版本号"""
    version = sys.version_info
    return f"{version.major}.{version.minor}"

def get_torch_version():
    """获取Torch完整版本号"""
    try:
        import torch
        version_parts = torch.__version__.split('.')[:2]  # 取前两部分
        return f"torch{'.'.join(version_parts)}"
    except ImportError:
        return "torch2.7"  # 默认版本

def get_platform():
    """获取系统平台"""
    system = platform.system().lower()
    if system == "windows":
        return "win_amd64"
    elif system == "linux":
        return "linux_x86_64"
    else:
        raise RuntimeError(f"Unsupported platform: {system}")

def _download_progress(count, block_size, total_size):
    """下载进度回调函数"""
    percent = int(count * block_size * 100 / total_size)
    print(f"\r下载进度: {percent}%", end="", flush=True)

def download_and_install(url, python_path=None):
    """下载并安装whl文件
    :param python_path: 指定的Python解释器路径，如果None则使用sys.executable
    """
    print(f"正在下载: {url}")
    try:
        script_dir = os.path.dirname(os.path.abspath(__file__))
        whl_name = os.path.basename(url)
        save_path = os.path.join(script_dir, whl_name)
        
        # 获取文件大小
        with urllib.request.urlopen(url) as response:
            file_size = int(response.headers['Content-Length'])
        
        print(f"文件大小: {file_size/1024/1024:.2f} MB")
        urllib.request.urlretrieve(url, save_path, reporthook=_download_progress)
        print("\n下载完成")
        print("下载完成，正在安装...")
        
        # 使用指定的Python路径或默认路径
        python_exec = python_path if python_path else sys.executable
        subprocess.check_call([python_exec, "-m", "pip", "install", save_path])
        print("安装成功!")
        return True
    except Exception as e:
        print(f"安装失败: {str(e)}")
        return False
    finally:
        if 'save_path' in locals() and os.path.exists(save_path):
            os.unlink(save_path)

def check_dependencies():
    """检查所有必需依赖是否可用"""
    required = ['sys', 'platform', 'subprocess', 'urllib', 'tempfile', 'os']
    missing = []
    for lib in required:
        try:
            __import__(lib)
        except ImportError:
            missing.append(lib)
    
    if missing:
        print("错误: 缺少以下依赖库:")
        for lib in missing:
            print(f"  - {lib}")
        print(f"\n请使用以下命令安装: {sys.executable} -m pip install {' '.join(missing)}")
        return False
    return True

def main():
    if not check_dependencies():
        sys.exit(1)
    import argparse
    parser = argparse.ArgumentParser(description='安装匹配的nunchaku wheel文件')
    parser.add_argument('--python-path', help='指定Python解释器路径')
    args = parser.parse_args()

    try:
        # 获取系统信息
        py_version = get_python_version()
        torch_version = get_torch_version()
        platform_suffix = get_platform()
        
        print(f"检测到环境: Python {py_version}, {torch_version}, {platform_suffix}")

        # 检查Python版本是否支持
        if py_version not in VERSION_MAP:
            raise RuntimeError(f"不支持的Python版本: {py_version}")
        # 构建whl文件名
        base_url = "https://github.com/mit-han-lab/nunchaku/releases/download/v0.3.1dev20250607"
        whl_name = f"nunchaku-0.3.1.dev20250607+{torch_version}-{VERSION_MAP[py_version]}-{platform_suffix}.whl"
        download_url = f"{base_url}/{whl_name}"

        print(f"尝试安装: {whl_name}")
        
        # 下载并安装
        if not download_and_install(download_url, args.python_path):
            raise RuntimeError("安装失败")

    except Exception as e:
        print(f"错误: {str(e)}")
        print("请手动下载适合您系统的whl文件:")
        print("支持的版本组合:")
        for py_ver in VERSION_MAP:
            for torch_ver in ["torch2.5", "torch2.6", "torch2.7"]:
                print(f"- Python {py_ver} + {torch_ver}")
        return 1

    return 0

if __name__ == "__main__":
    sys.exit(main())